{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9fd54a32",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/jerryjliu/llama_index/blob/main/docs/examples/llm/fireworks_cookbook.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e3a8796-edc8-43f2-94ad-fe4fb20d70ed",
   "metadata": {},
   "source": [
    "# Fireworks Function Calling Cookbook\n",
    "\n",
    "Fireworks.ai supports function calling for its LLMs, similar to OpenAI. This lets users directly describe the set of tools/functions available and have the model dynamically pick the right function calls to invoke, without complex prompting on the user's part.\n",
    "\n",
    "Since our Fireworks LLM directly subclasses OpenAI, we can use our existing abstractions with Fireworks.\n",
    "\n",
    "We show this on three levels: directly on the model API, as part of a Pydantic Program (structured output extraction), and as part of an agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f6f8702",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-fireworks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83ea30ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b070abb8-fa3f-4892-b23e-3ae91d0bf340",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"FIREWORKS_API_KEY\"] = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5497a17f-1099-4baf-884a-3620705be350",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from llama_index.llms.fireworks import Fireworks\n",
    "\n",
    "## define fireworks model\n",
    "llm = Fireworks(\n",
    "    model=\"accounts/fireworks/models/firefunction-v1\", temperature=0\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b007403c-6b7a-420c-92f1-4171d05ed9bb",
   "metadata": {},
   "source": [
    "## Function Calling on the LLM Module\n",
    "\n",
    "You can directly input function calls on the LLM module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "015c2d39",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ChatCompletionMessageToolCall(id='call_34ZaM0xPl1cveODjVUpO78ra', function=Function(arguments='{\"name\": \"Crazy in Love\", \"artist\": \"Beyonce\"}', name='Song'), type='function', index=0)]\n"
     ]
    }
   ],
   "source": [
    "from pydantic import BaseModel\n",
    "from llama_index.llms.openai.utils import to_openai_tool\n",
    "\n",
    "\n",
    "class Song(BaseModel):\n",
    "    \"\"\"A song with name and artist\"\"\"\n",
    "\n",
    "    name: str\n",
    "    artist: str\n",
    "\n",
    "\n",
    "# this converts pydantic model into function to extract structured outputs\n",
    "song_fn = to_openai_tool(Song)\n",
    "\n",
    "\n",
    "response = llm.complete(\"Generate a song from Beyonce\", tools=[song_fn])\n",
    "tool_calls = response.additional_kwargs[\"tool_calls\"]\n",
    "print(tool_calls)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f3e9bbe-e1ee-4396-8e03-0bb6455761fa",
   "metadata": {},
   "source": [
    "## Using a Pydantic Program\n",
    "\n",
    "Our Pydantic programs allow structured output extraction into a Pydantic object. `OpenAIPydanticProgram` takes advantage of function calling for structured output extraction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1332454-6aff-464b-a428-e4c94bd24fb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.program.openai import OpenAIPydanticProgram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40de5e32-e015-419a-8faf-01ffd0e45222",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_template_str = \"Generate a song about {artist_name}\"\n",
    "program = OpenAIPydanticProgram.from_defaults(\n",
    "    output_cls=Song, prompt_template_str=prompt_template_str, llm=llm\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "247c9c74-86c3-41a6-9579-93db817557c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "output = program(artist_name=\"Eminem\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1ebc517-469e-49b0-9197-708cf34b8454",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Song(name='Rap God', artist='Eminem')"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af7d9e57-5fa6-43ed-93f1-36841b688289",
   "metadata": {},
   "source": [
    "## Using An OpenAI Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8347ca9-98fb-4f65-a644-dc50abeb39fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.agent.openai import OpenAIAgent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8340c355-5e05-4c90-a1f7-719111ad4cd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.tools import BaseTool, FunctionTool\n",
    "\n",
    "\n",
    "def multiply(a: int, b: int) -> int:\n",
    "    \"\"\"Multiple two integers and returns the result integer\"\"\"\n",
    "    return a * b\n",
    "\n",
    "\n",
    "multiply_tool = FunctionTool.from_defaults(fn=multiply)\n",
    "\n",
    "\n",
    "def add(a: int, b: int) -> int:\n",
    "    \"\"\"Add two integers and returns the result integer\"\"\"\n",
    "    return a + b\n",
    "\n",
    "\n",
    "add_tool = FunctionTool.from_defaults(fn=add)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f96ada4-f117-4dd4-b726-6c02a6093eae",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = OpenAIAgent.from_tools(\n",
    "    [multiply_tool, add_tool], llm=llm, verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29df7ebd-74e4-4cd9-aedf-a4e63bd28857",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Added user message to memory: What is (121 * 3) + 42?\n",
      "=== Calling Function ===\n",
      "Calling function: multiply with args: {\"a\": 121, \"b\": 3}\n",
      "Got output: 363\n",
      "========================\n",
      "\n",
      "=== Calling Function ===\n",
      "Calling function: add with args: {\"a\": 363, \"b\": 42}\n",
      "Got output: 405\n",
      "========================\n",
      "\n",
      "The result of (121 * 3) + 42 is 405.\n"
     ]
    }
   ],
   "source": [
    "response = agent.chat(\"What is (121 * 3) + 42?\")\n",
    "print(str(response))"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "llama_index_v3",
   "language": "python",
   "name": "llama_index_v3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
